package com.itttq.gen.utils;

import com.alibaba.druid.sql.SQLUtils;
import com.alibaba.druid.sql.ast.SQLExpr;
import com.alibaba.druid.sql.ast.statement.SQLColumnDefinition;
import com.alibaba.druid.sql.ast.statement.SQLSelectOrderByItem;
import com.alibaba.druid.sql.ast.statement.SQLTableElement;
import com.alibaba.druid.sql.dialect.mysql.ast.MySqlPrimaryKey;
import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlCreateTableStatement;
import com.alibaba.druid.util.JdbcConstants;
import com.google.common.collect.Lists;
import com.itttq.gen.common.constant.Constants;
import com.itttq.gen.common.constant.DbConstant;
import com.itttq.gen.generator.to.ColumnTo;
import com.itttq.gen.generator.to.TableTo;
import com.itttq.gen.model.po.ProjectPo;
import com.itttq.gen.model.to.ResultTo;
import org.apache.commons.lang3.StringUtils;
import org.springframework.util.Assert;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.sql.*;
import java.util.List;

/**
 * 数据库操作工具类
 *
 * @createTime: Jul 3, 2013 2:41:49 PM
 * @author: <a href="mailto:g2room@163.com">天涯</a>
 * @version: 0.1
 * @lastVersion: 0.1
 * @updateTime:
 * @updateAuthor: <a href="mailto:g2room@163.com">天涯</a>
 * @changesSum:
 */
public class DbUtil {
    private static final String dbUrl = "jdbc:mysql://%s:%s/%s?useAffectedRows=true&useUnicode=true&characterEncoding=utf-8&tinyInt1isBit=false&serverTimezone=Asia/Shanghai";
    private static final String dbDriver = "com.mysql.cj.jdbc.Driver";
    private Connection conn;


    public static DbUtil getDbUtil(ProjectPo project) {
        DbUtil dbUtil = null;
        try {
            dbUtil = new DbUtil(project);
        } catch (SQLException e) {
            e.printStackTrace();
        } catch (ClassNotFoundException e) {
            e.printStackTrace();
        }
        Assert.notNull(dbUtil, ResultTo.getCodeAndMsg(ResultTo.ResultCode.FAIL));
        return dbUtil;
    }

    public DbUtil(ProjectPo project) throws SQLException, ClassNotFoundException {
        // 1.加载驱动
        Class.forName(dbDriver);
        // 2.得到连接
        conn = DriverManager.getConnection(String.format(dbUrl, project.getDbHost(), project.getDbPort(), project.getDbName()),
                project.getDbUsername(), project.getDbPassword());
    }

    /**
     * 获取数据库连接对象，这里只获得，需要自己在使用后进行关闭。
     *
     * @return 数据库连接实例
     */
    public Connection getConn() {
        return conn;
    }

    /**
     * 从数据源中获取字段名称及字段类型
     *
     * @param table 要获取字段的数据表名称
     * @return
     */
    public List<ColumnTo> getColumns(TableTo table, DatabaseMetaData metaData) {
        List<ColumnTo> columns = Lists.newArrayList();
        ResultSet rs = null;
        try {
            rs = metaData.getColumns(conn.getCatalog(), conn.getSchema(), table.getName(), "");
            rs.first();
            setColumnInfo(table, columns, rs);
            while (rs.next()) {
                setColumnInfo(table, columns, rs);
            }
        } catch (SQLException e) {
            e.printStackTrace();
        } finally {
            try {
                if (null != rs) {
                    rs.close();
                }
            } catch (SQLException e) {
                e.printStackTrace();
            }
        }
        return columns;
    }

    /**
     * 设置列信息
     *
     * @param table
     * @param columns
     * @param rs：     // COLUMN_NAME：字段名称
     *                // DATA_TYPE：字段数据类型(对应java.sql.Types中的常量)
     *                // TYPE_NAME：字段数据类型名称
     * @throws SQLException
     */
    private void setColumnInfo(TableTo table, List<ColumnTo> columns, ResultSet rs) throws SQLException {
        String fieldName = StringUtil.hump4Underline(rs.getString("COLUMN_NAME").replaceAll(Constants.GRAVE_ACCENT, ""));
        columns.add(ColumnTo.builder().column(rs.getString("COLUMN_NAME").replaceAll(Constants.GRAVE_ACCENT, ""))
                .fieldName(fieldName)
                .javaType(DbConstant.MAP_JDBC_JAVA_TYPE.get(rs.getInt("DATA_TYPE")))
                .remark(rs.getString("REMARKS").replaceAll(Constants.APOSTROPHE, ""))
                .autoincrement(rs.getString("IS_AUTOINCREMENT")).build());
        if (fieldName.equals(table.getPrimaryKey())) {
            table.setPrimaryKeyJavaType(DbConstant.MAP_JDBC_JAVA_TYPE.get(rs.getInt("DATA_TYPE")));
        }
    }

    /**
     * 从数据库获取表信息
     *
     * @param project
     * @return
     */
    public static List<TableTo> getTableInfoFromDb(ProjectPo project) {
        DbUtil dbUtil = getDbUtil(project);
        // 记录表数据，key为表名，值为表注释
        ResultSet rs = null;
        Connection conn = dbUtil.getConn();
        Statement statement = null;
        ResultSet pkInfo = null;
        List<TableTo> tables = Lists.newArrayList();
        try {
            DatabaseMetaData metaData = conn.getMetaData();
            String querySql = "SHOW TABLE STATUS";
            statement = conn.createStatement();
            rs = statement.executeQuery(querySql);
            String comment;
            String tableName;
            while (rs.next()) {
                tableName = rs.getString("Name").replaceAll(Constants.GRAVE_ACCENT, "");
                comment = rs.getString("Comment");
                TableTo table = TableTo.builder().name(tableName)
                        .remark(GenUtil.replaceWrapSymbol((null == comment ? null : comment.replaceAll(Constants.APOSTROPHE, ""))))
                        .build();
                // 获取表的主键名字
                pkInfo = metaData.getPrimaryKeys(conn.getCatalog(), conn.getSchema(), tableName);
                // 设置主键，如果有多个主键，则去掉主键。目前只支持一个主键
                while (pkInfo.next()) {
                    setPrimaryKey(pkInfo.getString("COLUMN_NAME"), table);
                }
                table.setColumns(dbUtil.getColumns(table, conn.getMetaData()));
                tables.add(table);
            }
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            try {
                if (null != pkInfo) {
                    pkInfo.close();
                }
            } catch (SQLException e) {
                e.printStackTrace();
            }
            try {
                if (null != rs) {
                    rs.close();
                }
            } catch (SQLException e) {
                e.printStackTrace();
            }
            try {
                if (null != statement) {
                    statement.close();
                }
            } catch (SQLException e) {
                e.printStackTrace();
            }
        }
        dbUtil.closeConn();
        return tables;
    }

    private static void setPrimaryKey(String pk, TableTo table) {
        if (StringUtils.isNotEmpty(table.getPrimaryKey())) {
            table.setPrimaryKey(null);
            table.setPrimaryKeyJavaType(null);
        } else {
            table.setPrimaryKey(StringUtil.hump4Underline(pk));
        }
    }

    /**
     * 关闭数据库连接。
     */
    public void closeConn() {
        if (null != conn) {
            try {
                conn.close();
            } catch (SQLException e) {
                e.printStackTrace();
            }
        }
    }

    public static List<TableTo> getTableInfoFromSql(String sql) {
        List<TableTo> tables = Lists.newArrayList();
        SQLUtils.parseStatements(sql, JdbcConstants.MYSQL, true).forEach(stmt -> {
            // 此处只处理Create语句
            if (stmt instanceof MySqlCreateTableStatement) {
                // 由于选择的是Mysql，此处强转为 MySqlCreateTableStatement 类
                MySqlCreateTableStatement createTableStatement = (MySqlCreateTableStatement) stmt;
                SQLExpr comment = createTableStatement.getComment();
                TableTo tableTo = TableTo.builder().name(createTableStatement.getTableName().replaceAll(Constants.GRAVE_ACCENT, ""))
                        .remark(GenUtil.replaceWrapSymbol(null == comment ? null : comment.toString().replaceAll(Constants.APOSTROPHE, "")))
                        .build();
                // 获取列信息
                List<SQLTableElement> elements = createTableStatement.getTableElementList();
                List<ColumnTo> columnTos = Lists.newArrayList();
                elements.forEach(x -> {
                    if (x instanceof MySqlPrimaryKey) {
                        // TODO 处理主键，目前是在下面当做普通的字段处理了
                        MySqlPrimaryKey mySqlPrimaryKey = (MySqlPrimaryKey) x;
                        List<SQLSelectOrderByItem> columns = mySqlPrimaryKey.getColumns();
                        columns.forEach(col -> setPrimaryKey(col.getExpr().toString().replaceAll(Constants.GRAVE_ACCENT, ""), tableTo));
                    } else if (x instanceof SQLColumnDefinition) {
                        SQLColumnDefinition columnDefinition = (SQLColumnDefinition) x;
                        String col = columnDefinition.getName().getSimpleName().replaceAll(Constants.GRAVE_ACCENT, "");
                        SQLExpr colComment = columnDefinition.getComment();
                        columnTos.add(ColumnTo.builder().column(col)
                                .fieldName(StringUtil.hump4Underline(col))
                                .remark(GenUtil.replaceWrapSymbol(null == colComment ? null : colComment.toString().replaceAll(Constants.APOSTROPHE, "")))
                                .javaType(DbConstant.MAP_JDBC_JAVA_TYPE.get(columnDefinition.getDataType().jdbcType())).build());
                    }
                });
                tableTo.setColumns(columnTos);
                // TODO 重新遍历获取主键的JAVA类型
                columnTos.forEach(col -> {
                    if (col.getFieldName().equals(tableTo.getPrimaryKey())) {
                        tableTo.setPrimaryKeyJavaType(col.getJavaType());
                    }
                });

                tables.add(tableTo);
            }
        });
        return tables;
    }

    /**
     * 从文件中读取SQL语句并转换为String
     *
     * @param is
     * @return
     * @throws IOException
     */
    public static String getSqlStrFromFile(InputStream is) throws IOException {
        StringBuilder sqlSb = new StringBuilder();
        InputStreamReader isr = null;
        try {
            isr = new InputStreamReader(is);
            BufferedReader br = new BufferedReader(isr);
            String line;
            while ((line = br.readLine()) != null) {
                sqlSb.append(line + "\n");
            }
        } catch (IOException e) {
            throw e;
        } finally {
            if (null != isr) {
                try {
                    isr.close();
                } catch (IOException e) {
                    throw e;
                }
            }
            if (null != is) {
                try {
                    is.close();
                } catch (IOException e) {
                    throw e;
                }
            }
        }
        return sqlSb.toString();
    }
}